#define OPENCV 1

#include <YOLOTrackingThread.hxx>
#include "FramePool.hxx"
#include "Sequence.hxx"
#include "VideoThread.hxx"
#include "TrackedObject.hxx"

YOLOTrackingThread::YOLOTrackingThread()
{
    _frameBuffer = 0;
    _stopTrackFlag = false;
    _trackingFlag = false;
    _pauseFlag = false;
    _trackerInitedFlag = false;
    _slowFlag = false;
    _nextFlag = false;
    _stopFlag = false;

    _detector = 0;

    strcpy(_trackerType, "CSRT");

    _status = TR_STATUS_FREE;
}

YOLOTrackingThread::~YOLOTrackingThread()
{
    if (_frameBuffer)
    {
        delete[] _frameBuffer;
        _frameBuffer = 0;
    }

    if( _detector )
    {
        delete _detector;
        _detector = 0;
    }
}

void YOLOTrackingThread::initialize()
{
    int pageSize = _inputPool->getPageSize();
    _outputPool->setPageSize(pageSize);
    _outputPool->setPageCount(10);
    _outputPool->initialize();

    _frameBuffer = new char[pageSize];

#ifdef _WIN64
    _detector = new Detector("D:/Fish/yolo-obj.cfg","D:/Fish/yolo-obj_last.weights");
#else
    _detector = new Detector("config/FishUpperView.cfg","data/Fish.lnk/FishUpperView_final.weights");
#endif

}

void YOLOTrackingThread::run()
{
    while (1)
    {
        if (_stopFlag) break;

        if (_pauseFlag)
        {
            msleep(30);
            continue;
        }

        if (_slowFlag)
            msleep(100);

        if (_inputPool->popFrame(_frameBuffer))
        {
            int frameIndex = *(int*)_frameBuffer;

            static int lastFrameIndex = frameIndex;

            if ((frameIndex < lastFrameIndex) || (frameIndex - lastFrameIndex != 1))
            {
                printf("Error!!!\n");
            }
            lastFrameIndex = frameIndex;

            if (_trackingFlag)
            {
                cv::Mat img(_model->getVideoHeight(), _model->getVideoWidth(), CV_8UC3, _frameBuffer + sizeof(VideoFrameInfo));

                //_model->emptyTrackingObjects();
                std::vector<TrackedObject*>* objects = _model->getTrackingObjects();
                //objects->clear();

                if (!_trackerInitedFlag)
                {
                    std::vector<bbox_t> trackedBox = _detector->detect(img);
                    trackedBox = _detector->tracking_id(trackedBox);

                    for(int i = 0; i < trackedBox.size();i++)
                    {
                        bbox_t bboxt = trackedBox.at(i);

                        cv::Rect2d bbox;
                        bbox.x = bboxt.x;
                        bbox.y = bboxt.y;
                        bbox.width = bboxt.w;
                        bbox.height = bboxt.h;

                        TrackedObject* object = new TrackedObject();

                        object->setClassIndex(bboxt.obj_id);
                        object->setTemplateBox(bbox);
                        object->pushRect(frameIndex, bbox);
                        object->setEndFrameIndex(frameIndex);
                        object->setStartFrameIndex(frameIndex + 1 - bboxt.frames_counter);
                        object->setObjectId(bboxt.track_id);

                        objects->push_back(object);
                    }

                    _trackerInitedFlag = true;

                    _status = TR_STATUS_RUNNING;
                }
                else
                {
                    std::vector<bbox_t> trackedBox = _detector->detect(img);
                    trackedBox = _detector->tracking_id(trackedBox);

                    for(int i = 0; i < trackedBox.size();i++)
                    {
                        bbox_t bboxt = trackedBox.at(i);

                        cv::Rect2d bbox;
                        bbox.x = bboxt.x;
                        bbox.y = bboxt.y;
                        bbox.width = bboxt.w;
                        bbox.height = bboxt.h;

                        // find object
                        bool found = false;
                        for(int j = 0; j < objects->size(); j ++)
                        {
                            TrackedObject* tobj = objects->at(j);
                            if( tobj->getObjectId() == bboxt.track_id )
                            {
                                tobj->setEndFrameIndex(frameIndex);
                                tobj->pushRect(frameIndex,bbox);
                                found = true;
                                break;
                            }
                        }

                        if(found) continue;

                        TrackedObject* object = new TrackedObject();
                        object->setClassIndex(bboxt.obj_id);
                        object->setTemplateBox(bbox);
                        object->pushRect(frameIndex, bbox);
                        object->setEndFrameIndex(frameIndex);
                        object->setStartFrameIndex(frameIndex + 1 - bboxt.frames_counter);
                        object->setObjectId(bboxt.track_id);

                        objects->push_back(object);
                    }

                    if (_stopTrackFlag)
                    {
                        _status = TR_STATUS_STOPPED;
                        _trackingFlag = false;
                        _stopTrackFlag = false;
                    }
                }

//                for (unsigned i = 0; i<_tracker->getObjects().size(); i++)
//                {

//                    TrackedObject* tobj = objects->at(i);
//                    cv::Rect2d rect = _tracker->getObjects()[i];
//                    tobj->setEndFrameIndex(frameIndex);
//                    tobj->pushRect(frameIndex, rect);
//                    //cv::rectangle(img, , cv::Scalar(255, 255, 0), 3);
//                }
            }

            while (_outputPool->isFull())
                msleep(30);

            if (!_outputPool->pushFrame(_frameBuffer))
            {
                msleep(20);
            }
        }
        else
        {
            msleep(20);
        }


        if (_nextFlag)
        {
            _pauseFlag = true;
            _nextFlag = false;
        }
    }
}

void YOLOTrackingThread::startTrack()
{
    _trackingFlag = true;
}

void YOLOTrackingThread::stopTrack()
{
    _stopTrackFlag = true;

    while( this->isRunning() && ( _status == TR_STATUS_RUNNING ))
        msleep(500);

    //_tracker.release();

    _trackerInitedFlag = false;
}

void YOLOTrackingThread::pause()
{
    _pauseFlag = true;
}

void YOLOTrackingThread::resume()
{
    _pauseFlag = false;
}

void YOLOTrackingThread::slow()
{
    _slowFlag = true;
}

void YOLOTrackingThread::fast()
{
    _slowFlag = false;
}

void YOLOTrackingThread::stop()
{
    _stopFlag = true;
}

void YOLOTrackingThread::next()
{
    _nextFlag = true;
    _pauseFlag = false;
}
